from torch import nn
import torch


class SkipGramModel(nn.Module):
    def __init__(self, vocab_size, embedding_dims, batch_size):  # embedding_dims 自定义（一般是3或者4）
        super().__init__()
        self.vocab_size = vocab_size
        self.embedding_dims = embedding_dims

        # 中心词的词向量
        self.center_embedding = nn.Embedding(vocab_size, embedding_dims)

        # 上下文的词向量
        self.context_embedding = nn.Embedding(vocab_size, embedding_dims)

        self.sigmoid = nn.Sigmoid()
        self.fc = nn.Linear(batch_size, 1)

    def forward(self, centers, contexts):
        c_embed = self.center_embedding(centers)
        t_embed = self.context_embedding(contexts)
        # 通过不同内容的词向量相乘来计算其相似度。
        scores = c_embed @ t_embed.T

        return self.sigmoid(self.fc(scores))  # 最终的结果会映射到0~1


if __name__ == '__main__':
    inputs = torch.randint(0, 10, (20, 2))
    outputs = torch.randint(0, 2, (20,))
    model = SkipGramModel(10, 5, 20)
    predicts = model(inputs[:, 0], inputs[:, 1])

    criterion = nn.BCEWithLogitsLoss()  # 官方建议的损失函数
    loss = criterion(predicts.squeeze(-1), outputs.float())
    print(loss)
